// Copyright (c) 2022 Presto Labs Pte. Ltd.
// Author: daniel

#pragma once

#include <memory>
#include <string>
#include <unordered_map>

#include "coin2/exchange/base/order/rate_config.h"
#include "coin2/exchange/base/order/rate_tracker.h"

namespace coin2::exchange::bybit_common::api_util {

enum HttpMethod : int {
  GET = 0,
  POST = 1,
};

namespace detail {

class IPRateTracker {
  using ActionType = coin2::exchange::base::order::IRateTracker::ActionType;
  using FixedWindowRateTracker = coin2::exchange::base::order::FixedWindowRateTracker;

 public:
  explicit IPRateTracker(const coin2::exchange::order::RateLimitInfoFuturesBybit& cfg);
  bool HasEnoughRateLimit(int64_t ts, ActionType action_type);
  void Push(int32_t num_reqs, int64_t timestamp, ActionType action_type);
  void SetRateLimitInfo(int32_t remain, HttpMethod http_method);
  auto GetMaxRateUsedRatio(int64_t ts) -> double;

 private:
  std::unique_ptr<FixedWindowRateTracker> get_method_rate_tracker_ = {};   // for GET
  std::unique_ptr<FixedWindowRateTracker> post_method_rate_tracker_ = {};  // For POST & DELETE
};

class ApiRateTracker {
  using ActionType = coin2::exchange::base::order::IRateTracker::ActionType;
  using FixedWindowRateTracker = coin2::exchange::base::order::FixedWindowRateTracker;

 public:
  explicit ApiRateTracker(
      const coin2::exchange::order::RateLimitInfoFuturesBybit& cfg,
      const std::string& rate_limit_category);
  bool HasEnoughRateLimit(int64_t ts, const std::string& endpoint);
  void Push(int32_t num_reqs, int64_t ts, const std::string& endpoint);
  void SetRateLimitInfo(const std::string& endpoint, int32_t max, int32_t remain);
  auto GetMaxRateUsedRatio(int64_t ts, const std::string& endpoint) -> double;

 private:
  const std::string rate_limit_category_;
  coin2::exchange::order::RateLimitInfoItem default_api_rate_limit_;
  coin2::exchange::order::RateLimitInfoFuturesBybit all_rate_limit_;
  std::unordered_map<std::string, std::unique_ptr<FixedWindowRateTracker>> rate_tracker_map_;
};

}  // namespace detail

class BybitRateTracker : public coin2::exchange::base::order::IRateTracker {
  enum BanType : int {
    SPECIFIC_API = 0,
    GET_API = 1,
    POST_API = 2,
  };

 public:
  using ActionType = coin2::exchange::base::order::IRateTracker::ActionType;
  using HttpContextPtr = coin2::exchange::base::order::IRateTracker::HttpContextPtr;
  using OrderLoggerCB = coin2::exchange::base::order::IRateTracker::OrderLoggerCB;
  using RateLimitReportWriter = coin2::exchange::base::order::RateLimitReportWriter;

 public:
  explicit BybitRateTracker(
      const MarketExchangeApi& mea,
      const coin2::exchange::order::OrderConnectionConfig& cfg);

  void SetRateLimitInfo(int32_t rate_limit, int32_t rate_limit_remaining) override { CHECK(false); }
  void Push(int32_t num_reqs, int64_t timestamp, ActionType action_type) override { CHECK(false); }
  bool HasEnoughRateLimit(int64_t timestamp, ActionType action_type) override { CHECK(false); }
  void ParseRateLimitInfo(const HttpContextPtr&, ActionType) override {}
  void SetRateLimitInfoString(const google::protobuf::Message* message) override;
  void WriteRateLimitReportMessage(OrderLoggerCB callback) override;
  double GetMaxRateUsedRatio(int64_t timestamp, ActionType action_type) override;
  void UpdateReqRateInfo(const nlohmann::json& info);

  bool HasEnoughRateLimit(
      int64_t ts,
      ActionType action_type,
      const std::string& endpoint,
      const std::string& bind_ip);
  void Push(
      int32_t num_reqs,
      int64_t ts,
      ActionType action_type,
      const std::string& endpoint,
      const std::string& bind_ip);

  void SetSumbitOrderEndpoint(const std::string& endpoint) { submit_order_endpoint_ = endpoint; }
  void SetCancelOrderEndpoint(const std::string& endpoint) { cancel_order_endpoint_ = endpoint; }
  void SetHttpGetEndpoints(const std::vector<std::string>& http_get_endpoints) {
    http_get_endpoints_ = http_get_endpoints;
  }
  void SetHttpPostEndpoints(const std::vector<std::string>& http_post_endpoints) {
    http_post_endpoints_ = http_post_endpoints;
  }

 private:
  void Validate(
      const MarketExchangeApi& mea,
      const coin2::exchange::order::RateLimitInfoFuturesBybit& rate_limit_info) const;

  bool CheckRateLimit(
      int64_t ts,
      ActionType action_type,
      const std::string& endpoint,
      const std::string& bind_ip);
  bool IsBannedPeriod(int64_t timestamp, const std::string& endpoint, const std::string& bind_ip);
  void BanApiForAPeriod(
      BanType ban_type,
      const std::string& endpoint,
      const std::string& bind_ip,
      int64_t retry_after);

 private:
  std::unordered_map<std::string, int64_t> enable_time_by_ip_map_;

  coin2::exchange::order::RateLimitInfoFuturesBybit rate_limit_info_;
  std::unique_ptr<detail::ApiRateTracker> query_api_rate_tracker_ = {};
  std::unique_ptr<detail::ApiRateTracker> place_api_rate_tracker_ = {};
  std::unique_ptr<detail::ApiRateTracker> cancel_api_rate_tracker_ = {};
  std::unordered_map<std::string, std::unique_ptr<detail::IPRateTracker>> ip_rate_tracker_map_;
  std::unique_ptr<RateLimitReportWriter> rate_limit_report_writer_ = {};

  // for get max used ratio
  std::string submit_order_endpoint_;
  std::string cancel_order_endpoint_;

  std::vector<std::string> http_get_endpoints_;
  std::vector<std::string> http_post_endpoints_;
};

}  // namespace coin2::exchange::bybit_common::api_util
